import numpy as np
import time

# 尝试导入tqdm，如果没有则使用简单的替代品
try:
    from tqdm import tqdm
    has_tqdm = True
except ImportError:
    has_tqdm = False
    print("提示: 没有找到tqdm模块，将使用简单进度显示。可以通过运行 'pip install tqdm' 安装。")

from vehicle_model import DoubleIntegratorVehicle
from lqr_controller import LQRController
from trajectory_generator import TrajectoryGenerator

class Simulation:
    """
    仿真环境
    """
    def __init__(self, vehicle, controller, trajectory_generator):
        """
        初始化仿真环境
        vehicle: 车辆模型
        controller: 控制器
        trajectory_generator: 轨迹生成器
        """
        self.vehicle = vehicle
        self.controller = controller
        self.trajectory_generator = trajectory_generator
        
    def run_simulation(self, reference_trajectory, initial_state, time_points, debug=False):
        """
        运行仿真
        reference_trajectory: 参考轨迹
        initial_state: 初始状态
        time_points: 时间点
        debug: 是否打印调试信息
        """
        # 记录开始时间
        start_time = time.time()
        
        # 设置初始状态
        current_state = initial_state.copy()
        self.vehicle.reset(initial_state)
        
        # 记录状态历史
        states_history = [current_state.copy()]
        controls_history = []
        
        # 使用tqdm进度条或简单进度显示
        time_steps = time_points[:-1]
        total_steps = len(time_steps)
        
        print(f"\n开始仿真 ({total_steps}步)...")
        
        if has_tqdm:
            # 使用tqdm进度条
            for i, t in enumerate(tqdm(time_steps, desc="仿真进度")):
                self._simulation_step(i, t, reference_trajectory, current_state, states_history, controls_history, debug)
                current_state = states_history[-1].copy()
        else:
            # 使用简单进度显示
            update_interval = max(1, total_steps // 20)  # 显示大约20次更新
            for i, t in enumerate(time_steps):
                self._simulation_step(i, t, reference_trajectory, current_state, states_history, controls_history, debug)
                current_state = states_history[-1].copy()
                
                # 显示简单进度
                if i % update_interval == 0 or i == total_steps - 1:
                    progress = (i + 1) / total_steps * 100
                    print(f"进度: {progress:.1f}% ({i+1}/{total_steps})", end="\r")
        
        print("\n仿真步骤完成!")
        
        # 记录结束时间
        end_time = time.time()
        elapsed_time = end_time - start_time
        
        # 转换为numpy数组
        states_history = np.array(states_history)
        controls_history = np.array(controls_history)
        
        # 打印仿真信息
        print(f"仿真完成! 耗时: {elapsed_time:.2f}秒")
        print(f"仿真步数: {len(states_history)}")
        
        # 计算和打印最终位置误差
        final_ref_state = reference_trajectory[len(states_history)-1] if len(states_history) <= len(reference_trajectory) else reference_trajectory[-1]
        final_state = states_history[-1]
        final_pos_error = np.linalg.norm(final_ref_state[:2] - final_state[:2])
        print(f"最终位置误差: {final_pos_error:.4f}")
        
        # 返回结果
        return {
            'time_points': time_points,
            'reference_trajectory': reference_trajectory,
            'states_history': states_history,
            'controls_history': controls_history,
            'simulation_time': elapsed_time
        }
        
    def _simulation_step(self, i, t, reference_trajectory, current_state, states_history, controls_history, debug):
        """单个仿真步骤"""
        # 获取当前参考状态
        reference_state = reference_trajectory[i].copy()
        
        # 计算控制输入
        u = self.controller.compute_control_input(current_state, reference_state)
        controls_history.append(u.copy())
        
        # 更新车辆状态
        new_state = self.vehicle.step(u).copy()
        states_history.append(new_state.copy())
        
        # 打印调试信息
        if debug and (i < 5 or i % 50 == 0):
            print(f"\nStep {i} (t={t:.1f}s):")
            print(f"  参考状态: {reference_state}")
            print(f"  当前状态: {current_state}")
            print(f"  新状态: {new_state}")
            print(f"  控制输入: {u}")
            pos_error = np.linalg.norm(reference_state[:2] - new_state[:2])
            print(f"  位置误差: {pos_error:.4f}")
            
            # 检查状态值是否合理
            if np.any(np.isnan(new_state)) or np.any(np.isinf(new_state)):
                print("警告: 状态包含NaN或Inf值!")
                return False
        
        return True

    def reset(self, initial_state=None, trajectory_type='circle'):
        """
        重置仿真
        initial_state: 初始状态
        trajectory_type: 轨迹类型
        """
        # 获取默认时间点
        dt = self.vehicle.dt
        sim_time = 30.0  # 默认仿真时间
        time_points = np.arange(0, sim_time, dt)
        
        # 生成参考轨迹
        self.reference_trajectory = self.trajectory_generator.generate_trajectory(trajectory_type, time_points)
        
        # 如果没有提供初始状态，使用轨迹起点
        if initial_state is None:
            initial_state = self.reference_trajectory[0].copy()
            
        # 重置车辆状态
        self.vehicle.reset(initial_state)
        
        # 清空历史记录
        self.states_history = [self.vehicle.get_state().copy()]
        self.control_inputs = []
        
        print("\n仿真已重置:")
        print(f"轨迹类型: {trajectory_type}")
        print(f"初始状态: {initial_state}")
        print(f"参考轨迹点数: {len(self.reference_trajectory)}")
        print("参考轨迹前5点:")
        for i in range(min(5, len(self.reference_trajectory))):
            print(f"  {i}: {self.reference_trajectory[i]}")
        
    def get_results(self):
        """获取仿真结果"""
        return {
            'time_points': self.time_points,
            'states_history': self.states_history,
            'reference_trajectory': self.reference_trajectory,
            'control_inputs': self.control_inputs
        }
        
    def calculate_errors(self):
        """计算轨迹跟踪误差"""
        if len(self.states_history) == 0:
            return None
            
        # 确保长度匹配
        min_length = min(len(self.reference_trajectory), len(self.states_history))
        ref_traj = self.reference_trajectory[:min_length]
        states = self.states_history[:min_length]
        
        # 计算位置误差
        position_errors = ref_traj[:, :2] - states[:, :2]
        distance_errors = np.sqrt(np.sum(position_errors**2, axis=1))
        
        # 计算RMSE
        rmse = np.sqrt(np.mean(distance_errors**2))
        
        # 计算最大误差
        max_error = np.max(distance_errors)
        
        return {
            'distance_errors': distance_errors,
            'rmse': rmse,
            'max_error': max_error
        } 